import os
import numpy as np
# from PIL import Image

import cv2
from torch.utils.data import dataset, DataLoader
import torchvision

import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2

def read_images(image_path):
    if not os.path.exists(image_path):
        raise FileNotFoundError('{} not found'.format(image_path))

    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)
    img /= 255.0 #[0~1]

    return img

# def read_images(image_path):
#     img = Image.open(image_path)

#     return img


class AliProducts(dataset.Dataset):
    def __init__(self, file_list, opt, mode='train', transforms=None):
        super().__init__()

        self.size = opt.size

        self.mode = mode
        self.transforms = transforms

        self.img_names = []
        self.labels = []

        fp = open(file_list, 'r')
        lines = fp.readlines()
        for line in lines:
            line = line.rstrip('\n')
            line = line.strip()

            if self.mode == "train" or self.mode == "eval":
                [img_n, label] = line.split(' ')
                self.img_names.append(img_n)
                self.labels.append(int(label))
            elif self.mode == "test":
                img_n = line.split(' ')
                self.img_names.append(img_n)
        fp.close()

    def __getitem__(self, index):
        image_name = self.img_names[index]
        if self.mode != "test":
            labels = self.labels[index]

        image_np = read_images(image_name)

        if self.transforms == None:
            self.transforms = A.Compose([
                A.Resize(height=self.size, width=self.size, p=1.0),
                A.Normalize(max_pixel_value=1.0, p=1.0),
                ToTensorV2(p=1.0),
            ], p=1.0)

        image = self.transforms(**{'image' : image_np})

        if self.mode == "train" or self.mode == "eval":
            return image['image'], labels
        elif self.mode == "test":
            return image['image']

    def __len__(self):
        return len(self.img_names)